import os
import sys
import numpy as np
import onnx
import onnx_graphsurgeon as gs

# === 路径加入工程根目录 ===
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

# ======= 保存文件夹设置 =======
folder = 'split_merge_onnx'
os.makedirs(folder, exist_ok=True)
pyfile = os.path.join(folder, 'merge_onnx_graph.py')

# ========== 路径配置 ==========
BACKBONE_ONNX =  os.path.join(PROJECT_ROOT, 'engine/quantization_onnx/yolov11n_int8_ALLConcat.onnx')
NMS_ONNX      =  os.path.join(PROJECT_ROOT, 'onnx/yolov11n_nms.onnx')
MERGED_ONNX   =  os.path.join(PROJECT_ROOT, 'engine/quantization_onnx/yolov11n_nms_merge.onnx')

BACKBONE_OUTPUT_NAME = None   # 自动查找最后一个输出
MERGE_INSERT_NODE = "/Transpose" # nms子图插入起点节点名

print("载入主干(on), NMS(on)...")
backbone_graph = gs.import_onnx(onnx.load(BACKBONE_ONNX))
nms_graph = gs.import_onnx(onnx.load(NMS_ONNX))

# 1. 定位 NMS分支 "/Transpose"节点及其所有下游节点
transpose_nodes = [node for node in nms_graph.nodes if node.name == MERGE_INSERT_NODE]
assert len(transpose_nodes) == 1, f"找不到或多个名为 {MERGE_INSERT_NODE} 的节点"
transpose_node = transpose_nodes[0]

def collect_downstream_nodes(start_node):
    selected = set()
    queue = [start_node]
    while queue:
        node = queue.pop()
        node_id = id(node)
        if node_id in selected:
            continue
        selected.add(node_id)
        for out in node.outputs:
            for n in nms_graph.nodes:
                if out in n.inputs:
                    queue.append(n)
    return [n for n in nms_graph.nodes if id(n) in selected]

nms_subgraph_nodes = collect_downstream_nodes(transpose_node)
nms_subgraph_nodes = list(nms_subgraph_nodes)
print(f"NMS子图节点数量: {len(nms_subgraph_nodes)}")

# 用于后续查常量（可选，便于调试）
used_tensor_names = set()
for node in nms_subgraph_nodes:
    used_tensor_names.update([v.name for v in node.inputs])
    used_tensor_names.update([v.name for v in node.outputs])

# 2. 自动获取主干输出名
if BACKBONE_OUTPUT_NAME is None:
    if len(backbone_graph.outputs) != 1:
        raise RuntimeError("主干输出不只一个，请指定 BACKBONE_OUTPUT_NAME")
    BACKBONE_OUTPUT_NAME = backbone_graph.outputs[0].name
print("主干输出名:", BACKBONE_OUTPUT_NAME)

# 3. 替换NMS子图起点input引用为主干输出Variable对象
transpose_input_name = transpose_node.inputs[0].name
print("NMS子图插入点input名:", transpose_input_name)

# 获取主干输出的Variable对象（ONE TRUE OBJECT）
backbone_output_var = None
for x in backbone_graph.tensors().values():
    if x.name == BACKBONE_OUTPUT_NAME:
        backbone_output_var = x
        break
assert backbone_output_var is not None, f"未找到主干输出变量 {BACKBONE_OUTPUT_NAME}"

if transpose_input_name != BACKBONE_OUTPUT_NAME:
    print(f"替换NMS子图中所有 {transpose_input_name} 为主干输出{BACKBONE_OUTPUT_NAME}")
    for n in nms_subgraph_nodes:
        n.inputs = [backbone_output_var if inp.name == transpose_input_name else inp for inp in n.inputs]
else:
    print("子图输入正好接主干，不需改名")

# 4. 以NMS子图的outputs作为新outputs
nms_final_outputs = nms_graph.outputs
backbone_graph.outputs = nms_final_outputs

# 5. 合并节点
backbone_graph.nodes.extend(nms_subgraph_nodes)

# 6. 合并常量
backbone_tensor_names = set(backbone_graph.tensors())
for t in nms_graph.tensors().values():
    import onnx_graphsurgeon as gs  # 确保gs是作用域内
    if isinstance(t, gs.Constant) and t.name not in backbone_tensor_names:
        backbone_graph.tensors()[t.name] = t

# 7. 清理与保存
print("图合并完成，清理...")
backbone_graph.cleanup()
onnx.save(gs.export_onnx(backbone_graph), MERGED_ONNX)
print(f"合并模型已保存到 {MERGED_ONNX}")